@ddundo is there a way to set up the docker container to install with pip install pygem[cpu] -f https://download.pytorch.org/whl/cpu.html?

Yup! That should be changed in the Dockerfile, in line 23.

Btw, I tested your changes locally and it still installs the CUDA version of pytorch:

In [2]: torch.__version__
Out[2]: '2.2.2+cu121'

hmm, that's strange but good to know, thanks!

Having some trouble getting this to work. If you have any ideas @ddundo, I'm all ears!

I tried to do this a few weeks ago and it's possible but quite complicated with poetry.

But I was thinking... I suggested this only for Docker purposes. So, as far as I'm concerned, I would be happy to just make the necessary changes in the Dockerfile, rather than changing pyproject.toml. Would you agree with this?

Edit: to clarify a bit more... if a user uses a machine without a GPU/CUDA and tries to install PyGEM, pytorch should automatically recognise that a CPU-only version should be installed. It's only Docker that requires some more explicit control in this case.

Thanks for the suggestion. I think that it makes sense to change the Dockerfile and install with the CPU option of PyTorch for Docker purposes as you suggested.

However, I did like the idea of making a CPU install option available if we're able. There may be some instances where someone wants to do a lightweight install (even for our own testing) and being able to specify to use the CPU PyTorch option would be nice. But I haven't figured out how to do this successfully yet.